[perf] fix: guard total_tokens behind is_hybrid_model to restore pure-transformer perf - #3585
Merged
Merged
Conversation
…-transformer perf PackedSeqParams.__post_init__ (added by MCore PR #3373) runs expensive CUDA ops (torch.repeat_interleave, arange, cat) to build seq_idx every forward step when total_tokens is set. seq_idx is only consumed by Mamba/SSM layers, so passing total_tokens unconditionally for all models caused ~10% throughput regression on pure-transformer configs (e.g. GB300 70B LoRA). Gate total_tokens on config.is_hybrid_model in gpt_step, llava_step, and vlm_step. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Chen Cui <chcui@nvidia.com>
Contributor
Author
|
/ok to test b3888ff |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
LGTM. The gating logic is correct and consistent across all three step files. getattr with a False default is a safe fallback, and the existing test_packed_seq_utils.py tests already cover both the total_tokens-present and absent paths at the utility layer.
yaoyu-33
approved these changes
Apr 30, 2026
yaoyu-33
added a commit
to DAISY-gh/Megatron-Bridge
that referenced
this pull request
Apr 30, 2026
Resolved vlm_step.py conflict: kept THD packing logic and adopted upstream's perf optimization that gates total_tokens behind config.is_hybrid_model (NVIDIA-NeMo#3585). Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
vasunvidia
pushed a commit
to vasunvidia/Megatron-Bridge
that referenced
this pull request
Jun 10, 2026
…-transformer perf (NVIDIA-NeMo#3585) Signed-off-by: Chen Cui <chcui@nvidia.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR #3484 added `total_tokens` unconditionally to `PackedSeqParams` in `gpt_step`, `llava_step`, and `vlm_step`. This triggers
`PackedSeqParams.post_init` (introduced by MCore PR #3373) on every forward step for all models — `torch.cat`, `repeat_interleave`,
`arange` — for all models.
`seq_idx` is only consumed by Mamba/SSM layers. For pure-transformer models this work is wasted, causing a ~10% throughput regression on
GB300 70B LoRA.
Fix
Gate `total_tokens` on `TransformerConfig.is_hybrid_model` in `gpt_step`, `llava_step`, and `vlm_step`. `config` is already available
via `get_model_config(model)` in each function.
Test plan